
import os
import torch
from torchvision import transforms
from PIL import Image
from pathlib import Path
import glob
import numpy as np
from typing import List, Tuple, Optional

class ImageFolderNoClass(torch.utils.data.Dataset):
    """加载指定文件夹中的所有图片"""
    SUPPORTED_EXTENSIONS = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.webp'}
    def __init__(self, root, transform=None, recursive=False):
        self.root = Path(root)
        self.transform = transform
        
        if not self.root.exists():
            raise FileExistsError(f"路径不存在： {self.root}")
        
        self.img_paths = []
        for file in os.listdir(root):
            if file.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp')):
                self.img_paths.append(os.path.join(root, file))
        self.img_paths.sort()
        
    def _get_image_paths(self, recusive:bool) -> List[Path]:
        """获取所有图片路径"""
        image_paths = []
        pattern = "**/*" if recusive else "*"
        
        for file_path in self.root.glob(pattern):
            if file_path.is_file() and file_path.suffix.lower() in self.SUPPORTED_EXTENSIONS:
                image_paths.append(file_path)
                
        return sorted(image_paths)

    def __len__(self):
        return len(self.img_paths)

    def __getitem__(self, idx):
        try:
            image = Image.open(self.img_paths[idx]).convert('RGB')
        except Exception as e:
            print(f"无法加载图片{self.img_paths[idx]}:{e}")
            image = Image.new('RGB',(640,640),(0,0,0)) 
        
        if self.transform:
            image = self.transform(image)
        return image
    
    def get_image_path(self,idx):
        """获取指定索引的图片路径"""
        return self.img_paths[idx]
    
def create_dataset(image_folder, img_size=640, batch_size=32, num_workers=4):
    """创建数据加载器"""
    transform = transforms.Compose([
        transforms.Resize((img_size, img_size)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], 
                           std=[0.229, 0.224, 0.225])
    ])
    
    dataset = ImageFolderNoClass(root=image_folder, transform=transform)
    
    dataloader = torch.utils.data.DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=num_workers,
        pin_memory=torch.cuda.is_available()
    )
    
    return dataloader

class LabelLoader:
    """标签加载器"""
    
    @staticmethod
    def load_label(label_path: str) -> np.ndarray:
        """
        加载YOLO格式标签文件
        
        Args:
            label_path: 标签文件路径
            
        Returns:
            标签数组 [class_id, x_center, y_center, width, height]
        """
        if not os.path.exists(label_path) or os.path.getsize(label_path) == 0:
            return np.zeros((0, 5))
        
        try:
            labels = np.loadtxt(label_path).reshape(-1, 5)
            return labels
        except Exception as e:
            print(f"Warning: Failed to load label {label_path}: {e}")
            return np.zeros((0, 5))

class DatasetLoader:
    """数据集加载器"""
    
    def __init__(self, img_dir: str, label_dir: str, img_ext: str = '*.jpg'):
        self.img_dir = img_dir
        self.label_dir = label_dir
        self.img_ext = img_ext
        self._validate_dirs()
    
    def _validate_dirs(self):
        """验证目录是否存在"""
        if not os.path.exists(self.img_dir):
            raise FileNotFoundError(f"Image directory not found: {self.img_dir}")
        if not os.path.exists(self.label_dir):
            raise FileNotFoundError(f"Label directory not found: {self.label_dir}")
    
    def get_image_paths(self) -> List[str]:
        """获取所有图像路径"""
        pattern = os.path.join(self.img_dir, self.img_ext)
        return sorted(glob.glob(pattern))
    
    def get_label_path(self, img_path: str) -> str:
        """根据图像路径获取对应标签路径"""
        img_name = os.path.basename(img_path)
        label_name = os.path.splitext(img_name)[0] + '.txt'
        return os.path.join(self.label_dir, label_name)
    
    def load_image_label_pair(self, img_path: str) -> Tuple[str, np.ndarray]:
        """加载图像标签对"""
        label_path = self.get_label_path(img_path)
        labels = LabelLoader.load_label(label_path)
        return img_path, labels
    
# 使用示例
if __name__ == "__main__":
    image_folder = "ppq_quantize/datasets/coco128/images/train2017"
    dataloader = create_dataset(image_folder, batch_size=16)
    
    print(f"数据集大小: {len(dataloader.dataset)}")
    
    # 测试加载第一个批次
    for images in dataloader:
        print(f"批次形状: {images.shape}")
        break